import os
from PIL import Image
from torch.utils.data import Dataset
import cv2
import numpy as np
from torch.utils.data import DataLoader
from transform import Compose, Normalize, ConvertDataType, Resize, RandomFilp, CenterCrop, RandomCrop, Scale, RandomScale, Pad

class DataAug(object):
    def __init__(self, size=256):
        self.size = size
    def __call__(self, data, label):
        '''
        data = cv2.resize(data, (self.size, self.size), interpolation=cv2.INTER_LINEAR)
        label = cv2.resize(label, (self.size, self.size), interpolation=cv2.INTER_NEAREST)
        '''
        augment = Compose([ Resize(self.size),
                            RandomScale(),
                            RandomFilp(),
                            Pad(self.size, mean_val=[0.485, 0.456, 0.406]),
                            RandomCrop(self.size),
                            ConvertDataType(),
                            Normalize(0, 1)])
                            
        data, label = augment(data, label)
        
        return data, label

class BacteriaDataset(object):
    def __init__(self, image_folder, image_list_file, mode, transforms=None, shuffle=True):
        self.image_folder = image_folder
        self.mode = mode
        self.image_list_file = image_list_file
        self.transforms = transforms
        self.shuffle = shuffle

        self.data_list = self.read_list()
    
    def read_list(self):
        data_list = []
        if self.mode == 'train':
            with open(self.image_list_file) as infile:
                for line in infile:
                    #data_path = os.path.join(self.image_folder, line.split()[0])
                    #label_path = os.path.join(self.image_folder, line.split()[1])
                    data_path = self.image_folder + line.split()[0]
                    label_path = self.image_folder + line.split()[1]
                    # print(f'data_path: {data_path}  label_path: {label_path}')
                    data_list.append((data_path, label_path))
        elif self.mode == 'test':
            pass
        return data_list

    def preprocess(self, data, label):
        h, w, c = data.shape
        h_gt, w_gt = label.shape
        assert h == h_gt, "Error"
        assert w == w_gt, "Error"

        if self.transforms:
            data, label = self.transforms(data, label)

        #_label = label
        label = label[:, :, np.newaxis]
        label = label.transpose(2,0,1)
        data = data.transpose(2,0,1)
        #label1 = _label[np.newaxis, :, :]
        #label2 = label.transpose(2,0,1)
        #print(label)
        #print(label1)
        #print(label2)
        # print(type(data))
        # data = data.permute(0,3,1,2)

        return data, label

    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, idx):
        data_path, label_path = self.data_list[idx][0], self.data_list[idx][1]
        data = cv2.imread(data_path, cv2.IMREAD_COLOR)
        data = cv2.cvtColor(data, cv2.COLOR_BGR2RGB)
        label = cv2.imread(label_path, cv2.IMREAD_GRAYSCALE)
        # print(data.shape, label.shape)
        data, label = self.preprocess(data, label)
        return data, label

def main():
    dataaug = DataAug(256)
    dataset = BacteriaDataset(
        image_folder='./Bacteria/training_data/',
        image_list_file='./Bacteria/training_data/list.txt',
        mode='train',
        transforms=dataaug,
        shuffle=True
    )
    dataloader = DataLoader(dataset, shuffle=True, batch_size=16, num_workers=1, pin_memory=True)
    num_epoch = 1
    for epoch in range(1, num_epoch+1):
        print(f'Epoch [{epoch}/{num_epoch}]:')
        for idx, (data, label) in enumerate(dataloader):
            print(f'Iter {idx}. Data shape: {data.shape}, label shape: {label.shape}')

if __name__ == "__main__":
    main()